import sys
import json
from ultralytics import YOLO

def predict(image_path):

    # 加载模型
    model = YOLO("D:/project/GitPro/model/major-events/backend/springboot-demo/src/main/java/com/qiu/demo/model/best.pt")

    # 进行预测
    results = model([image_path], stream=True)

    res = []
    for result in results:
        if len(result.boxes.cls) > 0:
            for i in range(len(result.boxes.cls)):
                lb_id = int(result.boxes.cls[i].item())
                lb = result.names[lb_id]
                xsd = result.boxes.conf[i].item()
                zb = result.boxes.xyxy[i].tolist()
                res.append({'类别': lb, '相似度': xsd, '坐标': zb})
    return res
    # 返回 JSON 格式的结果
    return json.dumps(res, ensure_ascii=False)

if __name__ == '__main__':
    # if len(sys.argv) != 2:
    #     print("Usage: python yolo_predict.py <image_path>")
    #     sys.exit(1)

    # image_path = sys.argv[1]
    result = predict('D:/project/GitPro/model/major-events/backend/springboot-demo/src/main/resources/static/image/adl-29-cam0-rgb-020_png.rf.cf47c8de779b9a7eb7f0363deb22e39f.jpg')
    print(result)